import numpy as np
from collections import defaultdict

class Apriori:
    def __init__(self, min_support=0.2):
        self.min_support = min_support
        self.frequent_itemsets = []

    def generate_c1(self, dataset):
        return list(map(frozenset, [ [item] for item in set(np.concatenate(dataset)) ]))

    def scan_dataset(self, dataset, candidates):
        item_count = defaultdict(int)
        for trans in dataset:
            for cand in candidates:
                if cand.issubset(trans):
                    item_count[cand] += 1
        total = len(dataset)
        return [k for k, v in item_count.items() if v/total >= self.min_support], \
               {k: v/total for k, v in item_count.items()}

    def apriori_gen(self, Lk, k):
        return [a | b for i, a in enumerate(Lk) for b in Lk[i+1:] if list(a)[:k-2] == list(b)[:k-2]]

    def fit(self, dataset):
        C1 = self.generate_c1(dataset)
        D = list(map(set, dataset))
        L1, support = self.scan_dataset(D, C1)
        self.frequent_itemsets = [L1]
        k = 2
        while len(self.frequent_itemsets[k-2]) > 0:
            Ck = self.apriori_gen(self.frequent_itemsets[k-2], k)
            Lk, supK = self.scan_dataset(D, Ck)
            support.update(supK)
            self.frequent_itemsets.append(Lk)
            k += 1
        return self.frequent_itemsets, support

# 可直接使用（无需修改路径，基于内存数据集）
if __name__ == "__main__":
    # 购物篮数据集（示例）
    dataset = [
        ['牛奶', '面包', '鸡蛋'], ['牛奶', '面包', '火腿'], ['牛奶', '薯片'],
        ['面包', '鸡蛋'], ['牛奶', '面包', '鸡蛋', '火腿'], ['面包', '火腿'],
        ['薯片', '可乐'], ['牛奶', '薯片', '可乐'], ['面包', '薯片'], ['牛奶', '面包', '薯片']
    ]
    apriori = Apriori(min_support=0.3)
    itemsets, support = apriori.fit(dataset)
    for k, items in enumerate(itemsets, 1):
        if items:
            print(f"{k}-频繁项集：{[set(item) for item in items]}")